package cn.itwen.star.graph.core.service.impl;

import cn.hutool.core.util.StrUtil;
import cn.itwen.star.graph.comfyui.client.pojo.ComfyuiTask;
import cn.itwen.star.graph.core.service.RedisService;
import com.alibaba.fastjson2.JSON;
import org.redisson.api.RLock;
import org.redisson.api.RedissonClient;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.ZSetOperations;
import org.springframework.stereotype.Service;

import java.time.Duration;

@Service
public class redisServiceImpl implements RedisService {

    private final static String TASK_KEY_PREFIX = "task_";
    private final static String DISTRIBUTED_ID_KEY = "DISTRIBUTED_ID";
    private final static String DISTRIBUTED_QUEUE_KEY = "DISTRIBUTED_QUEUE";

    private final static String RUN_TASK_KEY = "run_task_";

    public final static String QUEUE_LOCK_KEY = "QUEUE_LOCK_KEY";

    @Autowired
    private StringRedisTemplate stringRedisTemplate;

    @Autowired
    private RedissonClient redissonClient;


    /**
     * 添加任务
     * @param comfyuiTask
     * @return
     */
    @Override
    public ComfyuiTask addQueueTask(ComfyuiTask comfyuiTask) {
        //1.获取分布式id
        Long idx = stringRedisTemplate.opsForValue().increment(DISTRIBUTED_ID_KEY);
        //设置回任务对象中
        comfyuiTask.setIndex(idx);
        //2.添加任务到redis的String中
        stringRedisTemplate.opsForValue().set(TASK_KEY_PREFIX + comfyuiTask.getId(), JSON.toJSONString(comfyuiTask));
        //3，任务id压入队列的zset中
        stringRedisTemplate.opsForZSet().add(DISTRIBUTED_QUEUE_KEY,comfyuiTask.getId(),idx);
        return comfyuiTask;
    }


    /**
     * 弹出序列号最小的任务
     * @return
     */
    @Override
    public ComfyuiTask popQueueTask() {
        Long size = stringRedisTemplate.opsForZSet().size(DISTRIBUTED_QUEUE_KEY);
        if(size>0) {
            ZSetOperations.TypedTuple<String> task = stringRedisTemplate.opsForZSet().popMin(DISTRIBUTED_QUEUE_KEY);
            if (task != null) {
                String value = task.getValue();
                String json = stringRedisTemplate.opsForValue().getAndDelete(String.format("%s%s", TASK_KEY_PREFIX, value));
                if (StrUtil.isNotEmpty(json)) {
                    ComfyuiTask comfyuiTask = JSON.parseObject(json, ComfyuiTask.class);
                    return comfyuiTask;
                }
            }
        }
        return null;
    }


    /**
     * 缓存comfyui任务id和任务对象
     * @param promptId
     * @param comfyuiTask
     */
    @Override
    public void addStartedTask(String promptId, ComfyuiTask comfyuiTask) {
        stringRedisTemplate.opsForValue().set(RUN_TASK_KEY + promptId, JSON.toJSONString(comfyuiTask), Duration.ofMinutes(10));
    }


    /**
     * 获取缓存的comfyui任务对象
     * @param promptId
     * @return
     */
    @Override
    public ComfyuiTask getStartedTask(String promptId) {
       String json = stringRedisTemplate.opsForValue().get(RUN_TASK_KEY + promptId);
       if (StrUtil.isNotEmpty(json)) {
           return JSON.parseObject(json, ComfyuiTask.class);
       }
       return null;
    }

    /**
     * 判断是否有任务
     * @return
     */
    @Override
    public boolean hasQueueTask() {
        return stringRedisTemplate.opsForZSet().size(DISTRIBUTED_QUEUE_KEY) > 0;
    }

    /**
     * 取消任务
     * @param tempId
     * @return
     */
    @Override
    public ComfyuiTask cancelTask(String tempId) {
        RLock lock = redissonClient.getLock(QUEUE_LOCK_KEY);
        if(lock.tryLock()){
            try{
                boolean isDel = stringRedisTemplate.opsForZSet().remove(DISTRIBUTED_QUEUE_KEY, tempId)>0;
                if(isDel){
                    String json = stringRedisTemplate.opsForValue().getAndDelete(String.format("%s%s", TASK_KEY_PREFIX, tempId));
                    if(StrUtil.isNotEmpty(json)){
                        ComfyuiTask comfyuiTask = JSON.parseObject(json, ComfyuiTask.class);
                        return comfyuiTask;
                    }
                }
            }finally {
                lock.unlock();
            }
        }
        return null;
    }

    /**
     * 返回任务
     * @param comfyuiTask
     * @param index
     */
    @Override
    public void returnTask(ComfyuiTask comfyuiTask, Long index) {
        // 添加任务数据
        stringRedisTemplate.opsForValue().set(String.format("%s%s",TASK_KEY_PREFIX,comfyuiTask.getId()),JSON.toJSONString(comfyuiTask));
        // 压入队列
        stringRedisTemplate.opsForZSet().add(DISTRIBUTED_QUEUE_KEY,comfyuiTask.getId(),index);
    }
}
